#include "constraint_parser.h"

static std::string label_type[] = {"global", "local", "heap", "deref", "offset"};

int mypointsto::global_label::label_cmp(label *rhs){
	if(type<rhs->type)return -1;
	else if(type>rhs->type) return 1;
	else{
		global_label *gl = (global_label *) rhs;
		if(this->global_v<gl->global_v)return -1;
		else if(this->global_v==gl->global_v)return 0;
		else return 1;
	}
}

int mypointsto::local_label::label_cmp(label *rhs){
	if(type<rhs->type)return -1;
	else if(type>rhs->type) return 1;
	else{
		local_label *ll = (local_label*) rhs;
		if(func<ll->func)return -1;
		else if(func>ll->func)return 1;
		else{
			if(value<ll->value)return -1;
			else if(value>ll->value) return 1;
			else return 0;
		}
	}
}

int mypointsto::heap_label::label_cmp(label *rhs){
	if(type<rhs->type)return -1;
	else if(type>rhs->type)return 1;
	else{
		heap_label *hl = (heap_label*) rhs;
		if(alloc_inst<hl->alloc_inst)return -1;
		else if(alloc_inst>hl->alloc_inst)return 1;
		else return 0;
	}
}

int mypointsto::deref_label::label_cmp(label *rhs){
	if(type<rhs->type)return -1;
	else if(type>rhs->type) return 1;
	else{
		deref_label *dl = (deref_label*) rhs;
		int ret = lab->label_cmp(dl->lab);
		return ret;
	}
}

mypointsto::label *mypointsto::constraint_builder::build_value_label(llvm::Value *v){
	if(llvm::isa<llvm::GlobalVariable>(v)){
		llvm::GlobalVariable *gv = llvm::dyn_cast<llvm::GlobalVariable>(v);
		return new global_label(gv);
	}else if(llvm::isa<llvm::AllocaInst>(v)){
		llvm::AllocaInst *alloc_inst = llvm::dyn_cast<llvm::AllocaInst>(v);
		return new local_label(alloc_inst->getParent()->getParent(), v);
	}else if(llvm::isa<llvm::Argument>(v)){
		llvm::Argument *arg = llvm::dyn_cast<llvm::Argument>(v);
		return new local_label(arg->getParent(), v);
	}else if(llvm::isa<llvm::CastInst>(v)){
		llvm::CastInst *cast_inst = llvm::dyn_cast<llvm::CastInst>(v);
		return build_value_label(cast_inst->getOperand(0));
	}else if(llvm::isa<llvm::CallInst>(v)){
		llvm::CallInst *call_inst = llvm::dyn_cast<llvm::CallInst>(v);
		llvm::Function *called_func = call_inst->getCalledFunction();
		if(called_func && called_func->getNameStr()=="malloc")
			return new heap_label(call_inst);
		return new local_label(call_inst->getParent()->getParent(), call_inst);
	}else if(llvm::isa<llvm::LoadInst>(v)){
		llvm::LoadInst *load_inst = llvm::dyn_cast<llvm::LoadInst>(v);
		llvm::Value *addr = load_inst->getPointerOperand();
		label *tmp = build_value_label(addr);
		if(!tmp)return NULL;
		return new deref_label(tmp);
	}else if(llvm::isa<llvm::GetElementPtrInst>(v)){
		llvm::GetElementPtrInst *gep_inst = llvm::dyn_cast<llvm::GetElementPtrInst>(v);
		return build_value_label(gep_inst->getOperand(0));
	}else if(llvm::isa<llvm::ConstantExpr>(v)){
		llvm::ConstantExpr *cons_expr = llvm::dyn_cast<llvm::ConstantExpr>(v);
		if(cons_expr->getOpcode()>=29 && cons_expr->getOpcode()<=42)
			return build_value_label(cons_expr->getOperand(0));
	}
	return NULL;
}

mypointsto::label *mypointsto::constraint_builder::build_address_label(llvm::Value *v){
	mypointsto::label *ret=NULL;
	ret = mypointsto::constraint_builder::build_value_label(v);
	if(!ret)return NULL;
	return new deref_label(ret);
}

void mypointsto::constraint_builder::parse_address_label(mypointsto::label *l, bool &deref, int &label_id){
	if(l->type==deref_t && !is_var_value(l)){
		deref = true;
		label *l_base = ((deref_label*) l)->lab;
		bool temp;
		parse_value_label(l_base, temp, label_id);
	}else
		parse_value_label(l, deref, label_id);
}

void mypointsto::constraint_builder::parse_value_label(mypointsto::label *l, bool &deref, int &label_id){
	deref = false;
	bool temp = false;
	int tmp_id;
	mypointsto::deref_label *dl;
	switch (l->type){
		case deref_t:
			dl = (mypointsto::deref_label*) l;
			label_id = get_id(dl, true);
			if(!is_mem(dl->lab)){
				parse_value_label(dl->lab, temp, tmp_id);
				constraints.push_back(new load_constraint(label_id, tmp_id));
			}
			break;
		default: 
			label_id = get_id(l, true);
			if(is_mem(l)){
				tmp_id = get_id(new deref_label(l), true);
				constraints.push_back(new base_constraint(label_id, tmp_id));
			}
			break;
	}
}

bool mypointsto::constraint_builder::is_mem(mypointsto::label *l){
	switch (l->type){
		case global_t: return true;
		case local_t: return llvm::isa<llvm::AllocaInst>(((mypointsto::local_label*)l)->value);
		case heap_t: return true;
		default: return false;
	}
}

bool mypointsto::constraint_builder::is_var_value(mypointsto::label *l){
	if(l->type==deref_t){
		mypointsto::deref_label *dl = (mypointsto::deref_label*) l;
		return is_mem(dl->lab);
	}else
		return false;
}

void mypointsto::constraint_builder::for_each_function(){
	for(llvm::Module::iterator it = module->begin(); it!=module->end(); it++){
		llvm::Function *f = llvm::dyn_cast<llvm::Function>(it);
		if(!f->isDeclaration()){
			for(llvm::inst_iterator iter = inst_begin(f); iter != inst_end(f); iter++){
				llvm::Instruction *inst = iter.getInstructionIterator();
				populate_labels(inst);
			}
		}
	}
}

void mypointsto::constraint_builder::populate_labels(llvm::Instruction *inst){
	if(llvm::isa<llvm::StoreInst>(inst)){
		llvm::StoreInst *store_inst = llvm::dyn_cast<llvm::StoreInst>(inst);
		mypointsto::label *val = mypointsto::constraint_builder::build_value_label(store_inst->getValueOperand());
		mypointsto::label *addr = mypointsto::constraint_builder::build_address_label(store_inst->getPointerOperand());
		if(addr && val){
			bool deref1, deref2;
			int src_id, dst_id;
			parse_address_label(addr, deref1, dst_id);
			if(deref1){
				parse_value_label(val, deref2, src_id);
				deref2 = false;
			}else
				parse_address_label(val, deref2, src_id);
			if(deref1)
				constraints.push_back(new store_constraint(src_id, dst_id));
			else if(deref2)
				constraints.push_back(new load_constraint(src_id, dst_id));
			else
				constraints.push_back(new simple_constraint(src_id, dst_id));
		}
	}
}

void mypointsto::constraint_builder::dump_constraints(std::ostream &out){
	out<<labels.size()<<" labels processed and "<<constraints.size()<<" constraints resolved"<<std::endl;
	for(std::vector<mypointsto::constraint*>::const_iterator it = constraints.begin(); it!=constraints.end(); it++){
		print_constraint(out, *it);
		out<<std::endl;
	}
}

void mypointsto::constraint_builder::dump_labels(std::ostream &out){
	for(std::vector<label*>::iterator iter = labels.begin(); iter != labels.end(); iter++){
		print_label(out, *iter, true);
		out<<std::endl;
	}
}

void mypointsto::constraint_builder::print_constraint(std::ostream &out, mypointsto::constraint *c) {
	mypointsto::base_constraint *bc = NULL;
	mypointsto::simple_constraint *sc = NULL;
	mypointsto::load_constraint *lc = NULL;
	mypointsto::store_constraint *stc = NULL;
	switch(c->type){
		case base_t:
			out<<"base: (";
			bc = (mypointsto::base_constraint*) c;
			print_label(out, get_label(bc->val_id), false);
			out<<") <- ";
			print_label(out, get_label(bc->ptr_id), false);
			break;
		case simple_t:
			out<<"simple: ";
			sc = (mypointsto::simple_constraint*) c;
			print_label(out, get_label(sc->dst_id), false);
			out<<" = ";
			print_label(out, get_label(sc->src_id), false);
			break;
		case load_t:
			out<<"load: ";
			lc = (mypointsto::load_constraint*)c;
			print_label(out, get_label(lc->val_id), false);
			out<<" = *(";
			print_label(out, get_label(lc->addr_id), false);
			out<<")";
			break;
		case store_t:
			out<<"store: ";
			stc = (mypointsto::store_constraint*)c;
			out<<"*(";
			print_label(out, get_label(stc->addr_id), false);
			out<<") = ";
			print_label(out, get_label(stc->val_id), false);
			break;
		default:break;
	}
}

void mypointsto::constraint_builder::print_label(std::ostream &out, mypointsto::label* l, bool verbose) {
	if(!l){
		std::cout<<"null"<<std::endl;
		return;
	}
	mypointsto::global_label *gl = NULL;
	mypointsto::local_label *ll = NULL;
	mypointsto::heap_label *hl = NULL;
	mypointsto::deref_label *dl = NULL;
	std::string result, temp;
	llvm::raw_string_ostream raw_str_out(temp);
	int label_id = get_id(l, false);
	bool flag = (label_id < (int)labels.size());
	if(verbose)
		out<<label_type[l->type]<<" ";
	if(flag)
		out<<"v"<<label_id;
	if(!verbose)
		return;
	if(flag)
		out<<" ";
	out<<"[";
	switch(l->type){
		case global_t:
			gl = (mypointsto::global_label*) l;
			out<<gl->global_v->getNameStr();
			break;
		case local_t:
			ll = (mypointsto::local_label*) l;
			result = ll->func->getNameStr();
			result += ":";
			if(ll->value->hasName())
				result += ll->value->getNameStr();
			else{
				ll->value->print(raw_str_out);
				raw_str_out.flush();
				int i = 0;
				while (temp[i]==' ') i++;
				while ((temp[i]!=' ') && (temp[i]!='\0') && (temp[i]!='\n')) {
					result.push_back(temp[i]);
					i++;
				}
			}
			if(is_mem((mypointsto::label*)l))
				out<<"%";
			out<<result;
			break;
		case heap_t:
			hl = (mypointsto::heap_label*) l;
			out<<"heap alloc"; break;
		case deref_t:
			dl = (mypointsto::deref_label *) l;
			out<<"*(";
			print_label(out, dl->lab, true);
			out<<")";
			break;
		default: break;
	}
	out<<"]";
}

void mypointsto::constraint_builder::insert_label(mypointsto::label *l){
	if(!l){
		return;
	}
	if(get_id(l, false)>=0){
		return;
	}
	labels.push_back(l);
}

int mypointsto::constraint_builder::get_id(mypointsto::label *l, bool create){
	if(!l){
		return -1;
	}
	int idx = 0;
	for(std::vector<label*>::iterator it = labels.begin(); it != labels.end(); it++, idx++){
		if((*it)->label_cmp(l)==0){
			return idx;
		}
	}
	if(create && idx==(int)labels.size()){
		labels.push_back(l);
	}
	return idx;
}

mypointsto::label *mypointsto::constraint_builder::get_label(int id){
	if(id<(int)labels.size())
		return labels[id];
	return NULL;
}
